#!/usr/bin/env python
from MyPylike import MyLike

def mybool(Input):
    return {'True' : True, 'False' : False, 'T' : True, 'F' : False,
            't' : True, 'f' : False, 'TRUE' : True, 'FALSE' : False,
            "true" : True, "false" : False, "1" : True, "0" : False}.get(Input)

def RunBinnedLike(srcMaps, expCube, binnedExpMap, srcModel, irfs, centerSrc, emin=0., emax=1e10,
                  saveFolder='./', modelout1='model_pyLike_pass1.xml', modelout2='model_pyLike_pass2.xml',
                  TSLimit1=2., TSLimit2=25., tol1=0.001, tol2=1.e-5,
                  ULTSLimit=None, ULemin=None, ULemax=None,
                  npass=2, freezeWeak=False):
    'Main program of the cli'
    if npass < 1 or npass > 2: raise IOError('npass = 1|2')
    if not saveFolder:
        _saveFolder = ''
    elif saveFolder.endswith('/'):
        _saveFolder = saveFolder
    else:
        _saveFolder = saveFolder + '/'

    _modelout1 = _saveFolder + modelout1
    _modelout2 = _saveFolder + modelout2
    _resultdat = _saveFolder + 'results_pyLike.dat'
    _countsSpec = _saveFolder + 'counts_spectra_pyLike.fits'
    _covardat = _saveFolder + 'covariance_pyLike.dat'

    myBinned = MyLike.MyBinnedLikeObj(srcMaps, expCube, binnedExpMap, srcModel, irfs)
    if npass == 2:
        myBinned.initDRM(tol=tol1)
        myBinned.setERange(emin, emax)
        myBinned.saveEdisp()
        isGoodFit = myBinned.fitDRM(_modelout1, pokeSrc=centerSrc)

        if freezeWeak and isGoodFit:
            hasFreezed = myBinned.freezeWeak(centerSrc, TSLimit1, TSLimit2)
            if hasFreezed:
                myBinned.fitDRM(_modelout1, pokeSrc=centerSrc)

        myBinned.initMIN(tol=tol2, useBadFit=True)
        myBinned.restoreEdisp()
    else:
        myBinned.initMIN(tol=tol2)

    myBinned.setERange(emin, emax)
    myBinned.fitMIN(_modelout2)
    myBinned.paramsAtLimit()

    myBinned.writeResults(_resultdat, centerSrc, TSLimit=ULTSLimit, ULemin=ULemin, ULemax=ULemax)
    myBinned.writeCovar(_covardat)
    myBinned.writeCountsSpectra(_countsSpec)


def cli():
    import argparse

    parser = argparse.ArgumentParser(description='This is a binned likelihood analysis tools based on Fermi pyLikelihood package')
    parser.add_argument("srcMaps", type=str, help='Source Maps')
    parser.add_argument("expCube", type=str, help='Live Time Cube')
    parser.add_argument("binnedExpMap", type=str, help='Exposure Cube')
    parser.add_argument("srcModel", type=str, help='model.xml')
    parser.add_argument("irfs", type=str, help='IRFs')
    parser.add_argument("centerSrc", type=str, help='Center source')
    parser.add_argument("-emin", type=float, default=0, help='Use Max(emin_set, emin_data) as minimum energy in the fit')
    parser.add_argument("-emax", type=float, default=1e10, help='Use Min(emax_set, emax_data) as maximum energy in the fit')
    parser.add_argument("-saveFolder", type=str, default='./', help='Where to save output files')
    parser.add_argument("-modelout1", type=str, default='model_pyLike_pass1.xml', help='The output xml model of the first run')
    parser.add_argument("-modelout2", type=str, default='model_pyLike_pass2.xml', help='The output xml model of the second run')
    parser.add_argument("-TSLimit1", type=float, default=2., help='The sources whose TS lower than TSLimit1 is totally fixed in spectrum')
    parser.add_argument("-TSLimit2", type=float, default=25., help='The sources whose TS lower than TSLimit1 is fixed in index')
    parser.add_argument("-tol1", type=float, default=0.001, help='The tolerance of the first fit')
    parser.add_argument("-tol2", type=float, default=1.e-5, help='The tolerance of the second fit')
    parser.add_argument("-ULTSLimit", type=float, default=None, help='UpperLimit will be calculated when Center source whose TS lower than ULTSLimit')
    parser.add_argument("-ULemin", type=float, default=None, help='The emin[MeV] of the upper limit flux')
    parser.add_argument("-ULemax", type=float, default=None, help='The emax[MeV] of the upper limit flux')
    parser.add_argument("-npass", type=int, default=2, help='npass=1|2')
    parser.add_argument("-freezeWeak", type=mybool, default=False, help='Whether to freeze weak sources before fit in pass2')

    args = parser.parse_args()

    RunBinnedLike(args.srcMaps, args.expCube, args.binnedExpMap, args.srcModel, args.irfs, args.centerSrc,
                  args.emin, args.emax,
                  args.saveFolder, args.modelout1, args.modelout2,
                  args.TSLimit1, args.TSLimit2, args.tol1, args.tol2,
                  args.ULTSLimit, args.ULemin, args.ULemax,
                  args.npass, args.freezeWeak)


if __name__ == '__main__': cli()
